<!DOCTYPE html>

<html>
  <head>
    <meta charset="utf-8">
    
    <title>Writing custom array containers &mdash; NumPy v1.18 Manual</title>
    
    <link rel="stylesheet" type="text/css" href="../_static/css/spc-bootstrap.css">
    <link rel="stylesheet" type="text/css" href="../_static/css/spc-extend.css">
    <link rel="stylesheet" href="../_static/scipy.css" type="text/css" >
    <link rel="stylesheet" href="../_static/pygments.css" type="text/css" >
    <link rel="stylesheet" href="../_static/graphviz.css" type="text/css" >
    
    <script type="text/javascript">
      var DOCUMENTATION_OPTIONS = {
        URL_ROOT:    '../',
        VERSION:     '1.18.1',
        COLLAPSE_INDEX: false,
        FILE_SUFFIX: '.html',
        HAS_SOURCE:  false
      };
    </script>
    <script type="text/javascript" src="../_static/jquery.js"></script>
    <script type="text/javascript" src="../_static/underscore.js"></script>
    <script type="text/javascript" src="../_static/doctools.js"></script>
    <script type="text/javascript" src="../_static/language_data.js"></script>
    <script type="text/javascript" src="../_static/js/copybutton.js"></script>
    <link rel="author" title="About these documents" href="../about.html" >
    <link rel="index" title="Index" href="../genindex.html" >
    <link rel="search" title="Search" href="../search.html" >
    <link rel="top" title="NumPy v1.18 Manual" href="../index.html" >
    <link rel="up" title="NumPy basics" href="basics.html" >
    <link rel="next" title="Subclassing ndarray" href="basics.subclassing.html" >
    <link rel="prev" title="Structured arrays" href="basics.rec.html" > 
  </head>
  <body>
<div class="container">
  <div class="top-scipy-org-logo-header" style="background-color: #a2bae8;">
    <a href="../index.html">
      <img border=0 alt="NumPy" src="../_static/numpy_logo.png"></a>
    </div>
  </div>
</div>


    <div class="container">
      <div class="main">
        
	<div class="row-fluid">
	  <div class="span12">
	    <div class="spc-navbar">
              
    <ul class="nav nav-pills pull-left">
        <li class="active"><a href="https://numpy.org/">NumPy.org</a></li>
        <li class="active"><a href="https://numpy.org/doc">Docs</a></li>
        
        <li class="active"><a href="../index.html">NumPy v1.18 Manual</a></li>
        

          <li class="active"><a href="index.html" >NumPy User Guide</a></li>
          <li class="active"><a href="basics.html" accesskey="U">NumPy basics</a></li> 
    </ul>
              
              
    <ul class="nav nav-pills pull-right">
      <li class="active">
        <a href="../genindex.html" title="General Index"
           accesskey="I">index</a>
      </li>
      <li class="active">
        <a href="basics.subclassing.html" title="Subclassing ndarray"
           accesskey="N">next</a>
      </li>
      <li class="active">
        <a href="basics.rec.html" title="Structured arrays"
           accesskey="P">previous</a>
      </li>
    </ul>
              
	    </div>
	  </div>
	</div>
        

	<div class="row-fluid">
      <div class="spc-rightsidebar span3">
        <div class="sphinxsidebarwrapper">
  <h4>Previous topic</h4>
  <p class="topless"><a href="basics.rec.html"
                        title="previous chapter">Structured arrays</a></p>
  <h4>Next topic</h4>
  <p class="topless"><a href="basics.subclassing.html"
                        title="next chapter">Subclassing ndarray</a></p>
<div id="searchbox" style="display: none" role="search">
  <h4>Quick search</h4>
    <div>
    <form class="search" action="../search.html" method="get">
      <input type="text" style="width: inherit;" name="q" />
      <input type="submit" value="search" />
      <input type="hidden" name="check_keywords" value="yes" />
      <input type="hidden" name="area" value="default" />
    </form>
    </div>
</div>
<script type="text/javascript">$('#searchbox').show(0);</script>
        </div>
      </div>
          <div class="span9">
            
        <div class="bodywrapper">
          <div class="body" id="spc-section-body">
            
  <div class="section" id="module-numpy.doc.dispatch">
<span id="writing-custom-array-containers"></span><span id="basics-dispatch"></span><h1>Writing custom array containers<a class="headerlink" href="#module-numpy.doc.dispatch" title="Permalink to this headline">¶</a></h1>
<p id="dispatch-mechanism">Numpy’s dispatch mechanism, introduced in numpy version v1.16 is the
recommended approach for writing custom N-dimensional array containers that are
compatible with the numpy API and provide custom implementations of numpy
functionality. Applications include <a class="reference external" href="http://dask.pydata.org">dask</a> arrays, an
N-dimensional array distributed across multiple nodes, and <a class="reference external" href="https://docs-cupy.chainer.org/en/stable/">cupy</a> arrays, an N-dimensional array on
a GPU.</p>
<p>To get a feel for writing custom array containers, we’ll begin with a simple
example that has rather narrow utility but illustrates the concepts involved.</p>
<div class="doctest highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
<span class="gp">&gt;&gt;&gt; </span><span class="k">class</span> <span class="nc">DiagonalArray</span><span class="p">:</span>
<span class="gp">... </span>    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">N</span><span class="p">,</span> <span class="n">value</span><span class="p">):</span>
<span class="gp">... </span>        <span class="bp">self</span><span class="o">.</span><span class="n">_N</span> <span class="o">=</span> <span class="n">N</span>
<span class="gp">... </span>        <span class="bp">self</span><span class="o">.</span><span class="n">_i</span> <span class="o">=</span> <span class="n">value</span>
<span class="gp">... </span>    <span class="k">def</span> <span class="fm">__repr__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
<span class="gp">... </span>        <span class="k">return</span> <span class="sa">f</span><span class="s2">&quot;</span><span class="si">{self.__class__.__name__}</span><span class="s2">(N=</span><span class="si">{self._N}</span><span class="s2">, value=</span><span class="si">{self._i}</span><span class="s2">)&quot;</span>
<span class="gp">... </span>    <span class="k">def</span> <span class="nf">__array__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
<span class="gp">... </span>        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_i</span> <span class="o">*</span> <span class="n">np</span><span class="o">.</span><span class="n">eye</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">_N</span><span class="p">)</span>
<span class="gp">...</span>
</pre></div>
</div>
<p>Our custom array can be instantiated like:</p>
<div class="doctest highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="n">arr</span> <span class="o">=</span> <span class="n">DiagonalArray</span><span class="p">(</span><span class="mi">5</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">arr</span>
<span class="go">DiagonalArray(N=5, value=1)</span>
</pre></div>
</div>
<p>We can convert to a numpy array using <a class="reference internal" href="../reference/generated/numpy.array.html#numpy.array" title="numpy.array"><code class="xref py py-func docutils literal notranslate"><span class="pre">numpy.array</span></code></a> or
<a class="reference internal" href="../reference/generated/numpy.asarray.html#numpy.asarray" title="numpy.asarray"><code class="xref py py-func docutils literal notranslate"><span class="pre">numpy.asarray</span></code></a>, which will call its <code class="docutils literal notranslate"><span class="pre">__array__</span></code> method to obtain a
standard <code class="docutils literal notranslate"><span class="pre">numpy.ndarray</span></code>.</p>
<div class="doctest highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="n">np</span><span class="o">.</span><span class="n">asarray</span><span class="p">(</span><span class="n">arr</span><span class="p">)</span>
<span class="go">array([[1., 0., 0., 0., 0.],</span>
<span class="go">       [0., 1., 0., 0., 0.],</span>
<span class="go">       [0., 0., 1., 0., 0.],</span>
<span class="go">       [0., 0., 0., 1., 0.],</span>
<span class="go">       [0., 0., 0., 0., 1.]])</span>
</pre></div>
</div>
<p>If we operate on <code class="docutils literal notranslate"><span class="pre">arr</span></code> with a numpy function, numpy will again use the
<code class="docutils literal notranslate"><span class="pre">__array__</span></code> interface to convert it to an array and then apply the function
in the usual way.</p>
<div class="doctest highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="n">np</span><span class="o">.</span><span class="n">multiply</span><span class="p">(</span><span class="n">arr</span><span class="p">,</span> <span class="mi">2</span><span class="p">)</span>
<span class="go">array([[2., 0., 0., 0., 0.],</span>
<span class="go">       [0., 2., 0., 0., 0.],</span>
<span class="go">       [0., 0., 2., 0., 0.],</span>
<span class="go">       [0., 0., 0., 2., 0.],</span>
<span class="go">       [0., 0., 0., 0., 2.]])</span>
</pre></div>
</div>
<p>Notice that the return type is a standard <code class="docutils literal notranslate"><span class="pre">numpy.ndarray</span></code>.</p>
<div class="doctest highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="nb">type</span><span class="p">(</span><span class="n">arr</span><span class="p">)</span>
<span class="go">numpy.ndarray</span>
</pre></div>
</div>
<p>How can we pass our custom array type through this function? Numpy allows a
class to indicate that it would like to handle computations in a custom-defined
way through the interaces <code class="docutils literal notranslate"><span class="pre">__array_ufunc__</span></code> and <code class="docutils literal notranslate"><span class="pre">__array_function__</span></code>. Let’s
take one at a time, starting with <code class="docutils literal notranslate"><span class="pre">_array_ufunc__</span></code>. This method covers
<a class="reference internal" href="../reference/ufuncs.html#ufuncs"><span class="std std-ref">Universal functions (ufunc)</span></a>, a class of functions that includes, for example,
<a class="reference internal" href="../reference/generated/numpy.multiply.html#numpy.multiply" title="numpy.multiply"><code class="xref py py-func docutils literal notranslate"><span class="pre">numpy.multiply</span></code></a> and <a class="reference internal" href="../reference/generated/numpy.sin.html#numpy.sin" title="numpy.sin"><code class="xref py py-func docutils literal notranslate"><span class="pre">numpy.sin</span></code></a>.</p>
<p>The <code class="docutils literal notranslate"><span class="pre">__array_ufunc__</span></code> receives:</p>
<ul class="simple">
<li><p><code class="docutils literal notranslate"><span class="pre">ufunc</span></code>, a function like <code class="docutils literal notranslate"><span class="pre">numpy.multiply</span></code></p></li>
<li><p><code class="docutils literal notranslate"><span class="pre">method</span></code>, a string, differentiating between <code class="docutils literal notranslate"><span class="pre">numpy.multiply(...)</span></code> and
variants like <code class="docutils literal notranslate"><span class="pre">numpy.multiply.outer</span></code>, <code class="docutils literal notranslate"><span class="pre">numpy.multiply.accumulate</span></code>, and so
on.  For the common case, <code class="docutils literal notranslate"><span class="pre">numpy.multiply(...)</span></code>, <code class="docutils literal notranslate"><span class="pre">method</span> <span class="pre">==</span> <span class="pre">'__call__'</span></code>.</p></li>
<li><p><code class="docutils literal notranslate"><span class="pre">inputs</span></code>, which could be a mixture of different types</p></li>
<li><p><code class="docutils literal notranslate"><span class="pre">kwargs</span></code>, keyword arguments passed to the function</p></li>
</ul>
<p>For this example we will only handle the method <code class="docutils literal notranslate"><span class="pre">__call__</span></code>.</p>
<div class="doctest highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="kn">from</span> <span class="nn">numbers</span> <span class="kn">import</span> <span class="n">Number</span>
<span class="gp">&gt;&gt;&gt; </span><span class="k">class</span> <span class="nc">DiagonalArray</span><span class="p">:</span>
<span class="gp">... </span>    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">N</span><span class="p">,</span> <span class="n">value</span><span class="p">):</span>
<span class="gp">... </span>        <span class="bp">self</span><span class="o">.</span><span class="n">_N</span> <span class="o">=</span> <span class="n">N</span>
<span class="gp">... </span>        <span class="bp">self</span><span class="o">.</span><span class="n">_i</span> <span class="o">=</span> <span class="n">value</span>
<span class="gp">... </span>    <span class="k">def</span> <span class="fm">__repr__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
<span class="gp">... </span>        <span class="k">return</span> <span class="sa">f</span><span class="s2">&quot;</span><span class="si">{self.__class__.__name__}</span><span class="s2">(N=</span><span class="si">{self._N}</span><span class="s2">, value=</span><span class="si">{self._i}</span><span class="s2">)&quot;</span>
<span class="gp">... </span>    <span class="k">def</span> <span class="nf">__array__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
<span class="gp">... </span>        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_i</span> <span class="o">*</span> <span class="n">np</span><span class="o">.</span><span class="n">eye</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">_N</span><span class="p">)</span>
<span class="gp">... </span>    <span class="k">def</span> <span class="nf">__array_ufunc__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">ufunc</span><span class="p">,</span> <span class="n">method</span><span class="p">,</span> <span class="o">*</span><span class="n">inputs</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
<span class="gp">... </span>        <span class="k">if</span> <span class="n">method</span> <span class="o">==</span> <span class="s1">&#39;__call__&#39;</span><span class="p">:</span>
<span class="gp">... </span>            <span class="n">N</span> <span class="o">=</span> <span class="kc">None</span>
<span class="gp">... </span>            <span class="n">scalars</span> <span class="o">=</span> <span class="p">[]</span>
<span class="gp">... </span>            <span class="k">for</span> <span class="nb">input</span> <span class="ow">in</span> <span class="n">inputs</span><span class="p">:</span>
<span class="gp">... </span>                <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">Number</span><span class="p">):</span>
<span class="gp">... </span>                    <span class="n">scalars</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="nb">input</span><span class="p">)</span>
<span class="gp">... </span>                <span class="k">elif</span> <span class="nb">isinstance</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="vm">__class__</span><span class="p">):</span>
<span class="gp">... </span>                    <span class="n">scalars</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="nb">input</span><span class="o">.</span><span class="n">_i</span><span class="p">)</span>
<span class="gp">... </span>                    <span class="k">if</span> <span class="n">N</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
<span class="gp">... </span>                        <span class="k">if</span> <span class="n">N</span> <span class="o">!=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_N</span><span class="p">:</span>
<span class="gp">... </span>                            <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="s2">&quot;inconsistent sizes&quot;</span><span class="p">)</span>
<span class="gp">... </span>                    <span class="k">else</span><span class="p">:</span>
<span class="gp">... </span>                        <span class="n">N</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_N</span>
<span class="gp">... </span>                <span class="k">else</span><span class="p">:</span>
<span class="gp">... </span>                    <span class="k">return</span> <span class="bp">NotImplemented</span>
<span class="gp">... </span>            <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="vm">__class__</span><span class="p">(</span><span class="n">N</span><span class="p">,</span> <span class="n">ufunc</span><span class="p">(</span><span class="o">*</span><span class="n">scalars</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">))</span>
<span class="gp">... </span>        <span class="k">else</span><span class="p">:</span>
<span class="gp">... </span>            <span class="k">return</span> <span class="bp">NotImplemented</span>
<span class="gp">...</span>
</pre></div>
</div>
<p>Now our custom array type passes through numpy functions.</p>
<div class="doctest highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="n">arr</span> <span class="o">=</span> <span class="n">DiagonalArray</span><span class="p">(</span><span class="mi">5</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">np</span><span class="o">.</span><span class="n">multiply</span><span class="p">(</span><span class="n">arr</span><span class="p">,</span> <span class="mi">3</span><span class="p">)</span>
<span class="go">DiagonalArray(N=5, value=3)</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">np</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="n">arr</span><span class="p">,</span> <span class="mi">3</span><span class="p">)</span>
<span class="go">DiagonalArray(N=5, value=4)</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">np</span><span class="o">.</span><span class="n">sin</span><span class="p">(</span><span class="n">arr</span><span class="p">)</span>
<span class="go">DiagonalArray(N=5, value=0.8414709848078965)</span>
</pre></div>
</div>
<p>At this point <code class="docutils literal notranslate"><span class="pre">arr</span> <span class="pre">+</span> <span class="pre">3</span></code> does not work.</p>
<div class="doctest highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="n">arr</span> <span class="o">+</span> <span class="mi">3</span>
<span class="go">TypeError: unsupported operand type(s) for *: &#39;DiagonalArray&#39; and &#39;int&#39;</span>
</pre></div>
</div>
<p>To support it, we need to define the Python interfaces <code class="docutils literal notranslate"><span class="pre">__add__</span></code>, <code class="docutils literal notranslate"><span class="pre">__lt__</span></code>,
and so on to dispatch to the corresponding ufunc. We can achieve this
conveniently by inheriting from the mixin
<a class="reference internal" href="../reference/generated/numpy.lib.mixins.NDArrayOperatorsMixin.html#numpy.lib.mixins.NDArrayOperatorsMixin" title="numpy.lib.mixins.NDArrayOperatorsMixin"><code class="xref py py-class docutils literal notranslate"><span class="pre">NDArrayOperatorsMixin</span></code></a>.</p>
<div class="doctest highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="kn">import</span> <span class="nn">numpy.lib.mixins</span>
<span class="gp">&gt;&gt;&gt; </span><span class="k">class</span> <span class="nc">DiagonalArray</span><span class="p">(</span><span class="n">numpy</span><span class="o">.</span><span class="n">lib</span><span class="o">.</span><span class="n">mixins</span><span class="o">.</span><span class="n">NDArrayOperatorsMixin</span><span class="p">):</span>
<span class="gp">... </span>    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">N</span><span class="p">,</span> <span class="n">value</span><span class="p">):</span>
<span class="gp">... </span>        <span class="bp">self</span><span class="o">.</span><span class="n">_N</span> <span class="o">=</span> <span class="n">N</span>
<span class="gp">... </span>        <span class="bp">self</span><span class="o">.</span><span class="n">_i</span> <span class="o">=</span> <span class="n">value</span>
<span class="gp">... </span>    <span class="k">def</span> <span class="fm">__repr__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
<span class="gp">... </span>        <span class="k">return</span> <span class="sa">f</span><span class="s2">&quot;</span><span class="si">{self.__class__.__name__}</span><span class="s2">(N=</span><span class="si">{self._N}</span><span class="s2">, value=</span><span class="si">{self._i}</span><span class="s2">)&quot;</span>
<span class="gp">... </span>    <span class="k">def</span> <span class="nf">__array__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
<span class="gp">... </span>        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_i</span> <span class="o">*</span> <span class="n">np</span><span class="o">.</span><span class="n">eye</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">_N</span><span class="p">)</span>
<span class="gp">... </span>    <span class="k">def</span> <span class="nf">__array_ufunc__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">ufunc</span><span class="p">,</span> <span class="n">method</span><span class="p">,</span> <span class="o">*</span><span class="n">inputs</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
<span class="gp">... </span>        <span class="k">if</span> <span class="n">method</span> <span class="o">==</span> <span class="s1">&#39;__call__&#39;</span><span class="p">:</span>
<span class="gp">... </span>            <span class="n">N</span> <span class="o">=</span> <span class="kc">None</span>
<span class="gp">... </span>            <span class="n">scalars</span> <span class="o">=</span> <span class="p">[]</span>
<span class="gp">... </span>            <span class="k">for</span> <span class="nb">input</span> <span class="ow">in</span> <span class="n">inputs</span><span class="p">:</span>
<span class="gp">... </span>                <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">Number</span><span class="p">):</span>
<span class="gp">... </span>                    <span class="n">scalars</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="nb">input</span><span class="p">)</span>
<span class="gp">... </span>                <span class="k">elif</span> <span class="nb">isinstance</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="vm">__class__</span><span class="p">):</span>
<span class="gp">... </span>                    <span class="n">scalars</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="nb">input</span><span class="o">.</span><span class="n">_i</span><span class="p">)</span>
<span class="gp">... </span>                    <span class="k">if</span> <span class="n">N</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
<span class="gp">... </span>                        <span class="k">if</span> <span class="n">N</span> <span class="o">!=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_N</span><span class="p">:</span>
<span class="gp">... </span>                            <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="s2">&quot;inconsistent sizes&quot;</span><span class="p">)</span>
<span class="gp">... </span>                    <span class="k">else</span><span class="p">:</span>
<span class="gp">... </span>                        <span class="n">N</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_N</span>
<span class="gp">... </span>                <span class="k">else</span><span class="p">:</span>
<span class="gp">... </span>                    <span class="k">return</span> <span class="bp">NotImplemented</span>
<span class="gp">... </span>            <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="vm">__class__</span><span class="p">(</span><span class="n">N</span><span class="p">,</span> <span class="n">ufunc</span><span class="p">(</span><span class="o">*</span><span class="n">scalars</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">))</span>
<span class="gp">... </span>        <span class="k">else</span><span class="p">:</span>
<span class="gp">... </span>            <span class="k">return</span> <span class="bp">NotImplemented</span>
<span class="gp">...</span>
</pre></div>
</div>
<div class="doctest highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="n">arr</span> <span class="o">=</span> <span class="n">DiagonalArray</span><span class="p">(</span><span class="mi">5</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">arr</span> <span class="o">+</span> <span class="mi">3</span>
<span class="go">DiagonalArray(N=5, value=4)</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">arr</span> <span class="o">&gt;</span> <span class="mi">0</span>
<span class="go">DiagonalArray(N=5, value=True)</span>
</pre></div>
</div>
<p>Now let’s tackle <code class="docutils literal notranslate"><span class="pre">__array_function__</span></code>. We’ll create dict that maps numpy
functions to our custom variants.</p>
<div class="doctest highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="n">HANDLED_FUNCTIONS</span> <span class="o">=</span> <span class="p">{}</span>
<span class="gp">&gt;&gt;&gt; </span><span class="k">class</span> <span class="nc">DiagonalArray</span><span class="p">(</span><span class="n">numpy</span><span class="o">.</span><span class="n">lib</span><span class="o">.</span><span class="n">mixins</span><span class="o">.</span><span class="n">NDArrayOperatorsMixin</span><span class="p">):</span>
<span class="gp">... </span>    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">N</span><span class="p">,</span> <span class="n">value</span><span class="p">):</span>
<span class="gp">... </span>        <span class="bp">self</span><span class="o">.</span><span class="n">_N</span> <span class="o">=</span> <span class="n">N</span>
<span class="gp">... </span>        <span class="bp">self</span><span class="o">.</span><span class="n">_i</span> <span class="o">=</span> <span class="n">value</span>
<span class="gp">... </span>    <span class="k">def</span> <span class="fm">__repr__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
<span class="gp">... </span>        <span class="k">return</span> <span class="sa">f</span><span class="s2">&quot;</span><span class="si">{self.__class__.__name__}</span><span class="s2">(N=</span><span class="si">{self._N}</span><span class="s2">, value=</span><span class="si">{self._i}</span><span class="s2">)&quot;</span>
<span class="gp">... </span>    <span class="k">def</span> <span class="nf">__array__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
<span class="gp">... </span>        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_i</span> <span class="o">*</span> <span class="n">np</span><span class="o">.</span><span class="n">eye</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">_N</span><span class="p">)</span>
<span class="gp">... </span>    <span class="k">def</span> <span class="nf">__array_ufunc__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">ufunc</span><span class="p">,</span> <span class="n">method</span><span class="p">,</span> <span class="o">*</span><span class="n">inputs</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
<span class="gp">... </span>        <span class="k">if</span> <span class="n">method</span> <span class="o">==</span> <span class="s1">&#39;__call__&#39;</span><span class="p">:</span>
<span class="gp">... </span>            <span class="n">N</span> <span class="o">=</span> <span class="kc">None</span>
<span class="gp">... </span>            <span class="n">scalars</span> <span class="o">=</span> <span class="p">[]</span>
<span class="gp">... </span>            <span class="k">for</span> <span class="nb">input</span> <span class="ow">in</span> <span class="n">inputs</span><span class="p">:</span>
<span class="gp">... </span>                <span class="c1"># In this case we accept only scalar numbers or DiagonalArrays.</span>
<span class="gp">... </span>                <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">Number</span><span class="p">):</span>
<span class="gp">... </span>                    <span class="n">scalars</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="nb">input</span><span class="p">)</span>
<span class="gp">... </span>                <span class="k">elif</span> <span class="nb">isinstance</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="vm">__class__</span><span class="p">):</span>
<span class="gp">... </span>                    <span class="n">scalars</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="nb">input</span><span class="o">.</span><span class="n">_i</span><span class="p">)</span>
<span class="gp">... </span>                    <span class="k">if</span> <span class="n">N</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
<span class="gp">... </span>                        <span class="k">if</span> <span class="n">N</span> <span class="o">!=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_N</span><span class="p">:</span>
<span class="gp">... </span>                            <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="s2">&quot;inconsistent sizes&quot;</span><span class="p">)</span>
<span class="gp">... </span>                    <span class="k">else</span><span class="p">:</span>
<span class="gp">... </span>                        <span class="n">N</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_N</span>
<span class="gp">... </span>                <span class="k">else</span><span class="p">:</span>
<span class="gp">... </span>                    <span class="k">return</span> <span class="bp">NotImplemented</span>
<span class="gp">... </span>            <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="vm">__class__</span><span class="p">(</span><span class="n">N</span><span class="p">,</span> <span class="n">ufunc</span><span class="p">(</span><span class="o">*</span><span class="n">scalars</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">))</span>
<span class="gp">... </span>        <span class="k">else</span><span class="p">:</span>
<span class="gp">... </span>            <span class="k">return</span> <span class="bp">NotImplemented</span>
<span class="gp">... </span>   <span class="k">def</span> <span class="nf">__array_function__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">func</span><span class="p">,</span> <span class="n">types</span><span class="p">,</span> <span class="n">args</span><span class="p">,</span> <span class="n">kwargs</span><span class="p">):</span>
<span class="gp">... </span>       <span class="k">if</span> <span class="n">func</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">HANDLED_FUNCTIONS</span><span class="p">:</span>
<span class="gp">... </span>           <span class="k">return</span> <span class="bp">NotImplemented</span>
<span class="gp">... </span>       <span class="c1"># Note: this allows subclasses that don&#39;t override</span>
<span class="gp">... </span>       <span class="c1"># __array_function__ to handle DiagonalArray objects.</span>
<span class="gp">... </span>       <span class="k">if</span> <span class="ow">not</span> <span class="nb">all</span><span class="p">(</span><span class="nb">issubclass</span><span class="p">(</span><span class="n">t</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="vm">__class__</span><span class="p">)</span> <span class="k">for</span> <span class="n">t</span> <span class="ow">in</span> <span class="n">types</span><span class="p">):</span>
<span class="gp">... </span>           <span class="k">return</span> <span class="bp">NotImplemented</span>
<span class="gp">... </span>       <span class="k">return</span> <span class="n">HANDLED_FUNCTIONS</span><span class="p">[</span><span class="n">func</span><span class="p">](</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
<span class="gp">...</span>
</pre></div>
</div>
<p>A convenient pattern is to define a decorator <code class="docutils literal notranslate"><span class="pre">implements</span></code> that can be used
to add functions to <code class="docutils literal notranslate"><span class="pre">HANDLED_FUNCTIONS</span></code>.</p>
<div class="doctest highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="k">def</span> <span class="nf">implements</span><span class="p">(</span><span class="n">np_function</span><span class="p">):</span>
<span class="gp">... </span>   <span class="s2">&quot;Register an __array_function__ implementation for DiagonalArray objects.&quot;</span>
<span class="gp">... </span>   <span class="k">def</span> <span class="nf">decorator</span><span class="p">(</span><span class="n">func</span><span class="p">):</span>
<span class="gp">... </span>       <span class="n">HANDLED_FUNCTIONS</span><span class="p">[</span><span class="n">np_function</span><span class="p">]</span> <span class="o">=</span> <span class="n">func</span>
<span class="gp">... </span>       <span class="k">return</span> <span class="n">func</span>
<span class="gp">... </span>   <span class="k">return</span> <span class="n">decorator</span>
<span class="gp">...</span>
</pre></div>
</div>
<p>Now we write implementations of numpy functions for <code class="docutils literal notranslate"><span class="pre">DiagonalArray</span></code>.
For completeness, to support the usage <code class="docutils literal notranslate"><span class="pre">arr.sum()</span></code> add a method <code class="docutils literal notranslate"><span class="pre">sum</span></code> that
calls <code class="docutils literal notranslate"><span class="pre">numpy.sum(self)</span></code>, and the same for <code class="docutils literal notranslate"><span class="pre">mean</span></code>.</p>
<div class="doctest highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="nd">@implements</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">sum</span><span class="p">)</span>
<span class="gp">... </span><span class="k">def</span> <span class="nf">sum</span><span class="p">(</span><span class="n">arr</span><span class="p">):</span>
<span class="gp">... </span>    <span class="s2">&quot;Implementation of np.sum for DiagonalArray objects&quot;</span>
<span class="gp">... </span>    <span class="k">return</span> <span class="n">arr</span><span class="o">.</span><span class="n">_i</span> <span class="o">*</span> <span class="n">arr</span><span class="o">.</span><span class="n">_N</span>
<span class="gp">...</span>
<span class="gp">&gt;&gt;&gt; </span><span class="nd">@implements</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">mean</span><span class="p">)</span>
<span class="gp">... </span><span class="k">def</span> <span class="nf">mean</span><span class="p">(</span><span class="n">arr</span><span class="p">):</span>
<span class="gp">... </span>    <span class="s2">&quot;Implementation of np.mean for DiagonalArray objects&quot;</span>
<span class="gp">... </span>    <span class="k">return</span> <span class="n">arr</span><span class="o">.</span><span class="n">_i</span> <span class="o">/</span> <span class="n">arr</span><span class="o">.</span><span class="n">_N</span>
<span class="gp">...</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">arr</span> <span class="o">=</span> <span class="n">DiagonalArray</span><span class="p">(</span><span class="mi">5</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">np</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">arr</span><span class="p">)</span>
<span class="go">5</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">np</span><span class="o">.</span><span class="n">mean</span><span class="p">(</span><span class="n">arr</span><span class="p">)</span>
<span class="go">0.2</span>
</pre></div>
</div>
<p>If the user tries to use any numpy functions not included in
<code class="docutils literal notranslate"><span class="pre">HANDLED_FUNCTIONS</span></code>, a <code class="docutils literal notranslate"><span class="pre">TypeError</span></code> will be raised by numpy, indicating that
this operation is not supported. For example, concatenating two
<code class="docutils literal notranslate"><span class="pre">DiagonalArrays</span></code> does not produce another diagonal array, so it is not
supported.</p>
<div class="doctest highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="n">np</span><span class="o">.</span><span class="n">concatenate</span><span class="p">([</span><span class="n">arr</span><span class="p">,</span> <span class="n">arr</span><span class="p">])</span>
<span class="go">TypeError: no implementation found for &#39;numpy.concatenate&#39; on types that implement __array_function__: [&lt;class &#39;__main__.DiagonalArray&#39;&gt;]</span>
</pre></div>
</div>
<p>Additionally, our implementations of <code class="docutils literal notranslate"><span class="pre">sum</span></code> and <code class="docutils literal notranslate"><span class="pre">mean</span></code> do not accept the
optional arguments that numpy’s implementation does.</p>
<div class="doctest highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="n">np</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">arr</span><span class="p">,</span> <span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span>
<span class="go">TypeError: sum() got an unexpected keyword argument &#39;axis&#39;</span>
</pre></div>
</div>
<p>The user always has the option of converting to a normal <code class="docutils literal notranslate"><span class="pre">numpy.ndarray</span></code> with
<a class="reference internal" href="../reference/generated/numpy.asarray.html#numpy.asarray" title="numpy.asarray"><code class="xref py py-func docutils literal notranslate"><span class="pre">numpy.asarray</span></code></a> and using standard numpy from there.</p>
<div class="doctest highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="n">np</span><span class="o">.</span><span class="n">concatenate</span><span class="p">([</span><span class="n">np</span><span class="o">.</span><span class="n">asarray</span><span class="p">(</span><span class="n">arr</span><span class="p">),</span> <span class="n">np</span><span class="o">.</span><span class="n">asarray</span><span class="p">(</span><span class="n">arr</span><span class="p">)])</span>
<span class="go">array([[1., 0., 0., 0., 0.],</span>
<span class="go">       [0., 1., 0., 0., 0.],</span>
<span class="go">       [0., 0., 1., 0., 0.],</span>
<span class="go">       [0., 0., 0., 1., 0.],</span>
<span class="go">       [0., 0., 0., 0., 1.],</span>
<span class="go">       [1., 0., 0., 0., 0.],</span>
<span class="go">       [0., 1., 0., 0., 0.],</span>
<span class="go">       [0., 0., 1., 0., 0.],</span>
<span class="go">       [0., 0., 0., 1., 0.],</span>
<span class="go">       [0., 0., 0., 0., 1.]])</span>
</pre></div>
</div>
<p>Refer to the <a class="reference external" href="https://github.com/dask/dask">dask source code</a> and
<a class="reference external" href="https://github.com/cupy/cupy">cupy source code</a>  for more fully-worked
examples of custom array containers.</p>
<p>See also <a class="reference external" href="http://www.numpy.org/neps/nep-0018-array-function-protocol.html">NEP 18</a>.</p>
</div>


          </div>
        </div>
          </div>
        </div>
      </div>
    </div>

    <div class="container container-navbar-bottom">
      <div class="spc-navbar">
        
      </div>
    </div>
    <div class="container">
    <div class="footer">
    <div class="row-fluid">
    <ul class="inline pull-left">
      <li>
        &copy; Copyright 2008-2019, The SciPy community.
      </li>
      <li>
      Last updated on Feb 20, 2020.
      </li>
      <li>
      Created using <a href="http://sphinx.pocoo.org/">Sphinx</a> 2.4.2.
      </li>
    </ul>
    </div>
    </div>
    </div>
  </body>
</html>